using JLD2, Gadfly
using DataFrames

using FileIO
using Gadfly
using DataFrames

# First the noisy version **************************************************************************************************

include("cenf.jl")

cenf.readData("C:/Users/johannes.boehm/Dropbox/cenf/data/revision-2017")

println("Welfare counterfactual for Z_F (variable theta)")

f=load("../output_noise/ms_f_output_withvartheta_refined.jld2")
ms_fvals = f["ms_fvals"]
ms_θ = f["ms_θ"]
ms_β = f["ms_β"]
ms_η = f["ms_η"]
ms_γ = f["ms_γ"]
ms_logT = f["ms_logT"]
ms_logS = f["ms_logS"]

idx_best = findfirst(isequal(maximum(ms_fvals)), vec(ms_fvals));

bestθ = ms_θ[idx_best]
bestβ = vec(ms_β[idx_best,:])
bestη = vec(ms_η[idx_best,:])
bestγ = reshape(ms_γ[idx_best,:],35,35);
bestlogT = reshape(ms_logT[idx_best,:],109,35);
bestlogS = reshape(ms_logS[idx_best,:],109,35);

ioshares_before_VF = zeros(109,35,35);
ioshares_after_VF = zeros(109,35,35);
(changeArray_VF_noisy,ioshares_before_VF_noisy,ioshares_after_VF_noisy) = cenf.solveCountries(bestθ, bestβ, bestη, bestγ, 2.0, exp.(bestlogT), exp.(bestlogS),"","theta estimated, z^(2) used")




# Now the benchmark version **************************************************************************************************

include("cenf.jl")

cenf.readData("C:/Users/johannes.boehm/Dropbox/cenf/data/revision-2017")

println("Welfare counterfactual for Z_F (variable theta)")

f=load("../output/ms_f_output_withvartheta_refined.jld2")
ms_fvals = f["ms_fvals"]
ms_θ = f["ms_θ"]
ms_β = f["ms_β"]
ms_η = f["ms_η"]
ms_γ = f["ms_γ"]
ms_logT = f["ms_logT"]
ms_logS = f["ms_logS"]

idx_best = findfirst(isequal(maximum(ms_fvals)), vec(ms_fvals));

bestθ = ms_θ[idx_best]
bestβ = vec(ms_β[idx_best,:])
bestη = vec(ms_η[idx_best,:])
bestγ = reshape(ms_γ[idx_best,:],35,35);
bestlogT = reshape(ms_logT[idx_best,:],109,35);
bestlogS = reshape(ms_logS[idx_best,:],109,35);

ioshares_before_VF = zeros(109,35,35);
ioshares_after_VF = zeros(109,35,35);
(changeArray_VF,ioshares_before_VF,ioshares_after_VF) = cenf.solveCountries(bestθ, bestβ, bestη, bestγ, 2.0, exp.(bestlogT), exp.(bestlogS),"","theta estimated, z^(2) used")


# plot both
deltavec = cenf.δ[:,1,1]
Labels = cenf.countrynames
xticks = [0.25, 0.5, 0.75, 1]

# uncomment this when using Gadfly
df = DataFrame(x = deltavec, deltaU = 100.0 .* changeArray_VF[:,1], deltaU_noisy = 100.0 .* changeArray_VF_noisy[:,1], name = vec(cenf.countrynames))

dfPlot1 = DataFrame(delta = deltavec,dU = 100.0 .* changeArray_VF[:,1], Legend = "baseline")
dfPlot2 = DataFrame(delta = deltavec,dU = 100.0 .* changeArray_VF_noisy[:,1], Legend = "noisy z")
dfPlot = vcat(dfPlot1, dfPlot2)
p = Gadfly.plot(dfPlot, x=:delta, y=:dU , color="Legend", Geom.point, shape="Legend",
    Guide.xlabel("Enforcement cost"),
    Guide.ylabel("Counterfactual welfare increase, in percent"),
    Theme(default_color = colorant"red",
           highlight_width = 0pt)
     )
draw(SVG("../output_noise/welfare_us_variabletheta_f-noisy.svg",15cm, 10cm),p)
run(`cairosvg ../output_noise/welfare_us_variabletheta_f-noisy.svg -o ../output_noise/welfare_us_variabletheta_f-noisy.pdf`)